from langchain_community.chat_models import ChatOllama
from langchain_core.messages import HumanMessage
from langchain_core.prompts import ChatPromptTemplate


def compress_history(chat_history: list, max_length: int = 3) -> list:
    """
    压缩对话历史，保留最近对话并使用LLM摘要历史
    max_length: 保留的最近对话轮数
    """
    # 如果对话历史少于指定轮数，无需压缩
    if len(chat_history) <= max_length:
        return chat_history

    # 提取最近的对话
    recent_history = chat_history[-max_length:]

    # 摘要较早的历史
    summary = summarize_history(chat_history[:-max_length])

    # 返回摘要 + 最近对话
    return [{"role": "system", "content": f"历史摘要: {summary}"}] + recent_history


def summarize_history(history: list) -> str:
    """使用小型LLM生成历史摘要"""
    # 连接历史对话
    history_text = "\n".join(
        f"{'用户' if isinstance(msg, HumanMessage) else '助理'}: {msg.content}"
        for msg in history
    )

    # 使用本地小模型（可替换为其他模型）
    llm = ChatOllama(model="llama3:8b", temperature=0.2)

    # 摘要提示
    prompt = ChatPromptTemplate.from_template(
        "压缩以下对话历史，保留关键信息，50字以内:\n{history}\n摘要:"
    )
    chain = prompt | llm

    # 生成摘要
    result = chain.invoke({"history": history_text})
    return result.content